# coding=utf-8
__author__ = 'sdjkx'


import os
from PIL import Image
from libs.ustr import *
from aip import AipOcr

try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox,QFileDialog

class Ocr:
    def __init__(self):

        """ 你的 APPID AK SK """
        APP_ID = '23564473'
        API_KEY = 'Nw92zZFHZf2IBB5csVLiOZtj'
        SECRET_KEY = 'NXD2z4UHwmhMdbR7aKeInyouncO8qbiC'

        self.client = AipOcr(APP_ID, API_KEY, SECRET_KEY)
        # setConnectionTimeoutInMillis	建立连接的超时时间（单位：毫秒
        # setSocketTimeoutInMillis	通过打开的连接传输数据的超时时间（单位：毫秒）
        # self.client.setConnectionTimeoutInMillis(5000)
        # self.client.setSocketTimeoutInMillis(500)

    def ocrAllImage(self, labelImg, srcPathList):
        if srcPathList is None or len(srcPathList) <= 1:
            QMessageBox.information(labelImg, '错误', "必须指定源图片路径！")
            return
        else:
            self.srcPathList = srcPathList

        dirPath = ustr(QFileDialog.getExistingDirectory(labelImg,
                                                       "选择目标路径", "",  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirPath is None or len(dirPath) <= 1:
            QMessageBox.information(labelImg, '错误', "必须指定目标路径！")
            return
        else:
            self.targetPath = dirPath

        for path in self.srcPathList:
            self.ocr(path)

        QMessageBox.information(labelImg, '提示', "对文件执行文字识别完成！")

    def get_file_content(self, filePath):
        '''
        读取图片
        '''
        sImg = Image.open(filePath)
        w, h = sImg.size

        with open(filePath, 'rb') as fp:
            sImgBytes = fp.read()
        return w, h , sImgBytes

    def ocr(self, filePath = ""):
        fileName = filePath.split("\\")[-1]
        xmlPath = os.path.join(self.targetPath , fileName.split('.')[0] + ".xml")
        if os.path.exists(xmlPath):
            return

        print("OCR文件：", filePath)
        w, h, image = self.get_file_content(filePath)

        """ 调用通用文字识别, 图片参数为本地图片 """
        # self.client.basicGeneral(image);

        """ 如果有可选参数 """
        # options = {}
        # options["language_type"] = "CHN_ENG"  # 	识别语言类型，默认为CHN_ENG
        # options["detect_direction"] = "false"  # 是否检测图像朝向，默认不检测，即：false。
        #                                     # 朝向是指输入图像是正常方向、逆时针旋转90/180/270度。可选值包括:
        #                                     # - true：检测朝向；
        #                                     # - false：不检测朝向。
        # options["detect_language"] = "false"  # 是否检测语言，默认不检测。当前支持（中文、英语、日语、韩语）
        # options["probability"] = "false"       # 是否返回识别结果中每一行的置信度
        try:
            """ 带参数调用通用文字识别, 图片参数为本地图片 """
            response = self.client.general(image)
            print(response)
        except Exception as e:
            print(e)
            return

        if not ("words_result" in response):
            return

        xmlComment = "<annotation>\n" \
                     " <folder>train</folder>\n" \
                     " <filename>"+fileName+"</filename>\n" \
                     " <path>"+filePath+"</path>\n" \
                     " <source>\n" \
                     "     <database>Unknown</database>\n" \
                     " </source>\n" \
                     " <size>\n" \
                     "     <width>"+str(w)+"</width>\n" \
                     "     <height>"+str(h)+"</height>\n" \
                     "     <depth>3</depth>\n" \
                     " </size>\n" \
                     " <segmented>0</segmented>\n"

        words_result = response["words_result"]
        xmlTemp = ""
        rowid = 1
        for word in words_result:
            # probability = word["probability"]["average"]
            words = word["words"]
            location = word["location"]
            top = location["top"]
            left = location["left"]
            width = location["width"]
            height = location["height"]
            xmlTemp = xmlTemp + '   <object>\n' \
                                '       <name>'+words+'</name>\n' \
                                '       <pose>Unspecified</pose>\n' \
                                '       <truncated>0</truncated>\n' \
                                '       <difficult>0</difficult>\n' \
                                '       <rowid>'+str(rowid)+'</rowid>\n' \
                                '       <bndbox>\n' \
                                '           <xmin>'+str(left)+'</xmin>\n' \
                                '           <ymin>'+str(top)+'</ymin>\n' \
                                '           <xmax>'+str(left + width)+'</xmax>\n' \
                                '           <ymax>'+str(top + height)+'</ymax>\n' \
                                '       </bndbox>\n' \
                                '   </object>\n'
            rowid += 1
        xmlComment += xmlTemp + "</annotation>"

        self.saveXML(xmlComment, xmlPath)
        return

    def saveXML(self, xml, filePath):
        content=str(xml.encode('utf-8'),encoding='utf-8')

        # 前面省略，从下面直奔主题，举个代码例子：

        with open(filePath, 'a', encoding='utf-8') as file_handle:   # .txt可以不自己新建,代码会自动新建
            file_handle.write(content)     # 写入
            file_handle.write('\n')         # 有时放在循环里面需要自动转行，不然会覆盖上一条数据


